package hbase.template;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HConstants;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.PoolMap;
import org.springframework.beans.factory.DisposableBean;

import java.io.IOException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;

public class HBaseConnectionFactory implements DisposableBean {
    private Configuration conf;
    private volatile Connection connection;
    private volatile AsyncConnection asyncConnection;

    public HBaseConnectionFactory(Configuration conf) {
        this.conf = conf;
    }

    // Connection is thread safe
    public Connection getConnection() {
        if (connection == null) {
            synchronized (this) {
                if (connection == null) {
                    try {
                        connection = ConnectionFactory.createConnection(conf);
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                }
            }
        }
        return connection;
    }

    public AsyncConnection getAsyncConnection() {
        if (asyncConnection == null) {
            synchronized (this) {
                if (asyncConnection == null) {
                    try {
                        asyncConnection = ConnectionFactory.createAsyncConnection(conf).get();
                    } catch (InterruptedException | ExecutionException e) {
                        throw new RuntimeException(e);
                    }
                }
            }
        }
        return asyncConnection;
    }

    public Table getTable(String tableName) {
        try {
            return getConnection().getTable(TableName.valueOf(tableName));
        } catch (IOException e) {
            throw new RuntimeException("获取表失败：" + tableName, e);
        }
    }

    public Table getTable(String tableName, ExecutorService executorService) {
        try {
            return getConnection().getTable(TableName.valueOf(tableName), executorService);
        } catch (IOException e) {
            throw new RuntimeException("获取表失败：" + tableName, e);
        }
    }

    public AsyncTable<AdvancedScanResultConsumer> getAsyncTable(String tableName) {
        return getTableBuilder(tableName).build();
    }

    public AsyncTable<ScanResultConsumer> getAsyncTable(String tableName, ExecutorService pool) {
        return getTableBuilder(tableName, pool).build();
    }

    public AsyncTableBuilder<AdvancedScanResultConsumer> getTableBuilder(String tableName) {
        return getAsyncConnection().getTableBuilder(TableName.valueOf(tableName));
    }

    public AsyncTableBuilder<ScanResultConsumer> getTableBuilder(String tableName, ExecutorService pool) {
        return getAsyncConnection().getTableBuilder(TableName.valueOf(tableName), pool);
    }

    @Override
    public void destroy() throws Exception {
        if (connection != null) {
            connection.close();
        }
        if (asyncConnection != null) {
            asyncConnection.close();
        }
    }

    public void setIpcPoolType(PoolMap.PoolType ipcPoolType) {
        conf.set(HConstants.HBASE_CLIENT_IPC_POOL_TYPE, ipcPoolType.name());
    }

    public void setIpcPoolSize(int ipcPoolSize) {
        conf.set(HConstants.HBASE_CLIENT_IPC_POOL_SIZE, String.valueOf(ipcPoolSize));
    }

    public void setZookeeperQuorum(String quorum) {
        conf.set(HConstants.ZOOKEEPER_QUORUM, quorum.trim());
    }

    public void setZookeeperClientPort(int port) {
        conf.set(HConstants.ZOOKEEPER_CLIENT_PORT, String.valueOf(port));
    }

    /**
     * 设置写缓存大小，批量写入时不能超过写缓存大小。hbaseM默认2MB
     * @param sizeInMb 缓存大小，单位MB
     */
    public void setWriteBufferSize(int sizeInMb) {
        conf.set(ConnectionConfiguration.WRITE_BUFFER_SIZE_KEY, String.valueOf(sizeInMb * 1024L * 1024));
    }
}
